package cn.xiaopengstack.domain.agent.service.armory.busines.data.impl;

import cn.xiaopengstack.domain.agent.adapter.repository.IAgentRepository;
import cn.xiaopengstack.domain.agent.model.entity.ArmoryCommandEntity;
import cn.xiaopengstack.domain.agent.model.valobj.AiAgentEnumVO;
import cn.xiaopengstack.domain.agent.model.valobj.AiClientApiVO;
import cn.xiaopengstack.domain.agent.model.valobj.AiClientModelVO;
import cn.xiaopengstack.domain.agent.service.armory.busines.data.ILoadDataStrategy;
import cn.xiaopengstack.domain.agent.service.armory.factory.DefaultArmoryStrategyFactory;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ThreadPoolExecutor;

/**
 * @author jiangyangang
 */
@Service("aiClientModelLoadDataStrategy")
@Slf4j
public class AiClientModelLoadDataStrategy implements ILoadDataStrategy {
    @Resource
    private IAgentRepository agentRepository;
    @Resource
    private ThreadPoolExecutor threadPoolExecutor;

    @Override
    public void loadData(ArmoryCommandEntity armoryCommandEntity, DefaultArmoryStrategyFactory.DynamicContext dynamicContext) {
        List<String> modelIds = armoryCommandEntity.getCommandIdList();
        CompletableFuture<List<AiClientApiVO>> listCompletableFuture = CompletableFuture.supplyAsync(() -> {
            log.info("查询配置数据(ai_client_api) {}", modelIds);
            return agentRepository.queryAiClientApiVOListByModelIds(modelIds);
        }, threadPoolExecutor);

        CompletableFuture<List<AiClientModelVO>> aiClientModelListFuture = CompletableFuture.supplyAsync(() -> {
            log.info("查询配置数据(ai_client_model) {}", modelIds);
            return agentRepository.AiClientModelVOByModelIds(modelIds);
        }, threadPoolExecutor);

        CompletableFuture.allOf(listCompletableFuture).thenRun(() -> {
            dynamicContext.setValue(AiAgentEnumVO.AI_CLIENT_API.getDataName(), listCompletableFuture.join());
            dynamicContext.setValue(AiAgentEnumVO.AI_CLIENT_MODEL.getDataName(), aiClientModelListFuture.join());
        }).join();
    }
}
